import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
import threading
from queue import Queue
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args
from spirems.mod_helper import download_model
import argparse
from typing import Union
import platform
from copy import copy
import json
from rknn.api import RKNN
import uuid
from pathlib import Path


class ClipImgOnnx2RknnNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            **kwargs
        )

        self.platform = self.get_param("platform", "rk3588")
        self.onnx_model = self.get_param("onnx_model", "/home/orangepi/Downloads/clip_images.onnx")
        self.imgsz = self.get_param("imgsz", [224, 224])
        self.do_quant = self.get_param("do_quant", False)
        self.params_help()
        print("  platform:", self.platform)
        print("  onnx_model:", self.onnx_model)
        print("  imgsz:", self.imgsz)
        print("  do_quant:", self.do_quant)

        assert self.onnx_model.endswith(".onnx")

        # Create RKNN object
        rknn = RKNN(verbose=False)

        # Pre-process config
        print('--> Config model')
        rknn.config(target_platform=self.platform,
                    mean_values=[[0.48145466*255, 0.4578275*255, 0.40821073*255]],
                    std_values=[[0.26862954*255, 0.26130258*255, 0.27577711*255]])
        print('done')

        # Load model
        print('--> Loading model')
        ret = rknn.load_onnx(model=self.onnx_model,
                            inputs=['pixel_values'],
                            input_size_list=[[1, 3, self.imgsz[0], self.imgsz[1]]])
        if ret != 0:
            print('Load model failed!')
            exit(ret)
        print('done')

        # Build model
        print('--> Building model')
        ret = rknn.build(do_quantization=self.do_quant)
        if ret != 0:
            print('Build model failed!')
            exit(ret)
        print('done')

        # Export rknn model
        print('--> Export rknn model')
        output_path = self.onnx_model[:-4] + "rknn"
        ret = rknn.export_rknn(output_path)
        if ret != 0:
            print('Export rknn model failed!')
            exit(ret)
        print('done')

        # Release
        rknn.release()

        self.release()

    def release(self):
        BaseNode.release(self)

    def run(self):
        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = ClipImgOnnx2RknnNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
